import os
import glob
import argparse
from pathlib import Path
import gmalg

def verify(msk_s_file, mpk_s_file, sk_s_file, uid_file):
    hid_s = b"\x01"
    mpk_s_temp = read_binary_file_(mpk_s_file)
    mpk_s = bytes.fromhex("04") + mpk_s_temp
    msk_s = read_binary_file_(msk_s_file)
    sk_s_temp = read_binary_file_(sk_s_file)
    sk_s = bytes.fromhex("04") + sk_s_temp
    uid = read_binary_file_(uid_file)

    kgc = gmalg.SM9KGC(hid_s=hid_s, msk_s=msk_s, mpk_s=mpk_s)
    
    sk_s_creat = kgc.generate_sk_sign(uid)
    
    #print(sk_s.hex())
    #print(sk_s_creat.hex())
    assert sk_s == sk_s_creat, "验证失败" 
    
    sm9 = gmalg.SM9(hid_s=hid_s, mpk_s=mpk_s, sk_s=sk_s, uid=uid)
    
    message = b"Chinese IBS standard"
    h, s = sm9.sign(message)
    
    #print(h.hex())
    #print(S.hex())
    assert sm9.verify(message, h, s) == True, "验证失败"

def read_binary_file_(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        return binary_data


def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    parser.add_argument('mode', choices=['ok', 'speed'], default='ok', help='运行模式')
    args = parser.parse_args()
    directory = args.directory
    mode = args.mode

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    id_files = sorted(glob.glob(os.path.join(directory, 'id_*')))
    count = len(id_files)
    print(f"count={count}")

    # 创建输出内容
    merged_content = ""
    
    # 遍历0-9，合并对应的公钥和私钥文件
    for i in range(count):
        # 读取文件内容
        id_file = os.path.join(directory, f'id_{i}')
        ks_file = os.path.join(directory, f'ks_{i}')
        Ppub_s_file = os.path.join(directory, f'Ppub_s_{i}')
        dsA_file = os.path.join(directory, f'dsA_{i}')
        verify(ks_file, Ppub_s_file, dsA_file, id_file)

        id_content = read_binary_file(id_file)
        ks_content = read_binary_file(ks_file)
        Ppub_s_content = read_binary_file(Ppub_s_file)
        dsA_content = read_binary_file(dsA_file)

        # 添加到合并内容
        merged_content += f"用户标识= {id_content}\n"
        merged_content += f"用户主私钥= {ks_content}\n"
        merged_content += f"公钥= {Ppub_s_content}\n"
        merged_content += f"私钥= {dsA_content}\n"
        merged_content += "\n"
        
    # 写入到txt文件
    if mode == 'ok':
        output_file = 'sm9签名公钥.txt'
    else:
        output_file = 'sm9签名公钥生成速率.txt'


    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)
    
    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
